#include "Shader.h"
#include <cassert>
#include <d3dcompiler.h>

HRESULT Shader::CompileShaderFromFile(LPCWSTR szFileName, LPCSTR szEntryPoint, LPCSTR szShaderModel, ID3DBlob** ppBlobOut )
{
	HRESULT hr = S_OK;

	DWORD dwShaderFlags = D3DCOMPILE_ENABLE_STRICTNESS;
#if defined( DEBUG ) || defined( _DEBUG )
	// Set the D3DCOMPILE_DEBUG flag to embed debug information in the shaders.
	// Setting this flag improves the shader debugging experience, but still allows 
	// the shaders to be optimized and to run exactly the way they will run in 
	// the release configuration of this program.
	dwShaderFlags |= D3DCOMPILE_DEBUG;
	dwShaderFlags |= D3DCOMPILE_SKIP_OPTIMIZATION;
#endif

	ID3DBlob* pErrorBlob;
	hr = D3DX11CompileFromFile( szFileName, NULL, NULL, szEntryPoint, szShaderModel, 
		dwShaderFlags, 0, NULL, ppBlobOut, &pErrorBlob, NULL );
	if( FAILED(hr) )
	{
		if( pErrorBlob != NULL )
			OutputDebugStringA( (char*)pErrorBlob->GetBufferPointer() );
		if( pErrorBlob ) pErrorBlob->Release();
		return hr;
	}
	if( pErrorBlob ) pErrorBlob->Release();

	return S_OK;
}

HRESULT Shader::Init(
	ID3D11Device* device, 
	ID3D11DeviceContext* imDevice,
	const std::wstring& shaderName, 
	D3D11_PRIMITIVE_TOPOLOGY topology, 
	D3D11_INPUT_ELEMENT_DESC* layout,
	size_t nElements)
{
	m_device = device;
	m_imContext = imDevice;
	std::wstring vsPath = shaderName + L"_vs.hlsl"; 
	std::wstring psPath = shaderName + L"_ps.hlsl";
	m_topology = topology;

	HRESULT hr;
	//prepare vs
	ID3DBlob* pBlob = NULL;
	hr = CompileShaderFromFile(vsPath.c_str(), "main", "vs_5_0", &pBlob);
	assert(SUCCEEDED(hr));
	hr = device->CreateVertexShader(pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pVS);
	assert(SUCCEEDED(hr));
	//layout
	hr = device->CreateInputLayout(layout, (UINT)nElements, pBlob->GetBufferPointer(), pBlob->GetBufferSize(), &m_pLayout);
	pBlob->Release();
	pBlob = NULL;
	//prepare ps
	hr = CompileShaderFromFile(psPath.c_str(), "main", "ps_5_0", &pBlob);
	assert(SUCCEEDED(hr));
	hr = device->CreatePixelShader(pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pPS);
	assert(SUCCEEDED(hr));
	pBlob->Release();
	return hr;
}

HRESULT mdx::Shader::Init(const Param& param)
{
	m_device = param.Get<ID3D11Device*>("device");
	m_imContext = param.Get<ID3D11DeviceContext*>("imContext");
	std::wstring vsPath = param.Get<std::wstring>("shaderName") + L"_vs.hlsl"; 
	std::wstring psPath = param.Get<std::wstring>("shaderName") + L"_ps.hlsl";
	m_topology = param.Get<D3D11_PRIMITIVE_TOPOLOGY>("topology");

	HRESULT hr;
	//prepare vs
	ID3DBlob* pBlob = NULL;
	hr = CompileShaderFromFile(vsPath.c_str(), "main", "vs_5_0", &pBlob);
	assert(SUCCEEDED(hr));
	hr = m_device->CreateVertexShader(pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pVS);
	assert(SUCCEEDED(hr));
	//layout
	hr = m_device->CreateInputLayout(
		param.Get<D3D11_INPUT_ELEMENT_DESC*>("layout"), 
		param.Get<UINT>("nElements"),
		pBlob->GetBufferPointer(),
		pBlob->GetBufferSize(),
		&m_pLayout
		);
	pBlob->Release();
	pBlob = NULL;
	//prepare ps
	hr = CompileShaderFromFile(psPath.c_str(), "main", "ps_5_0", &pBlob);
	assert(SUCCEEDED(hr));
	hr = m_device->CreatePixelShader(pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pPS);
	assert(SUCCEEDED(hr));
	pBlob->Release();
	return hr;
}

Shader::Shader()
{
	ZeroMemory(this, sizeof(this));
}

Shader::~Shader()
{
	m_pLayout->Release();
	m_pVS->Release();
	m_pPS->Release();
	m_pVertexCBuffer->Release();
}

HRESULT Shader::Update(const void* pVertexCBuffer, const void* pGeometryCBuffer, const void* pFragmentCBuffer)
{
	//update will set shader as current
	assert(m_pVertexCBuffer);
	//set current
	m_imContext->VSSetShader(m_pVS, NULL, 0);
	m_imContext->VSSetConstantBuffers(0, 1, &m_pVertexCBuffer);
	m_imContext->PSSetShader(m_pPS, NULL, 0);
	m_imContext->PSSetConstantBuffers(0, 1, &m_pFragmentCBuffer);
	m_imContext->IASetInputLayout(m_pLayout);
	m_imContext->IASetPrimitiveTopology(m_topology);
	//update
	m_imContext->UpdateSubresource(m_pVertexCBuffer, 0 , NULL, pVertexCBuffer, 0, 0);
	return S_OK;
}

HRESULT Shader::InitCBuffer(size_t vertexCBSize, size_t fragmentCBSize)
{
	HRESULT hr;
	//create const buffer
	if(vertexCBSize != 0){
		D3D11_BUFFER_DESC bd;
		ZeroMemory(&bd, sizeof(bd));
		bd.Usage = D3D11_USAGE_DEFAULT;
		bd.ByteWidth = (UINT)vertexCBSize;
		bd.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
		hr = m_device->CreateBuffer(&bd, nullptr, &m_pVertexCBuffer);
		assert(SUCCEEDED(hr));
	}
	if(fragmentCBSize != 0){
		D3D11_BUFFER_DESC bd;
		ZeroMemory(&bd, sizeof(bd));
		bd.Usage = D3D11_USAGE_DEFAULT;
		bd.ByteWidth = (UINT)fragmentCBSize;
		bd.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
		hr = m_device->CreateBuffer(&bd, nullptr, &m_pFragmentCBuffer);
		assert(SUCCEEDED(hr));
	}
}
